from ..dlgo.rl.experience import ExperienceBuffer
import numpy as np 
import h5py

class ZeroExperienceCollector:
    def __init__(self) -> None:
        self.states = []
        self.visit_counts = []
        self.rewards = []
        self.current_episode_states = []
        self.current_episode_visit_counts = []

    def BeginEpisode(self):
        self.current_episode_states = []
        self.current_episode_visit_counts = []

    def RecordDecision(self, state, visit_counts):
        self.current_episode_states.append(state)
        self.current_episode_visit_counts.append(visit_counts)

    def CompleteEpisode(self, reward):
        num_states = len(self.current_episode_states)
        self.states += self.current_episode_states
        self.visit_counts += self.current_episode_visit_counts
        self.rewards += [reward for _ in range(num_states)]

        self.current_episode_states = []
        self.current_episode_visit_counts = []


class ZeroExperienceBuffer:
    def __init__(self, states, visit_counts, rewards) -> None:
        self.states = states
        self.visit_counts = visit_counts
        self.rewards = rewards

    def Serialize(self, f:h5py.File):
        e = 'experience'
        f.create_group(e)
        f[e].create_dataset('states', data=self.states)
        f[e].create_dataset('visit_counts', data=self.visit_counts)
        f[e].create_dataset('rewards', data=self.rewards)

def CombineExperiences(collectors):
    combined_states = np.concatenate([np.array(c.states) for c in collectors])
    combined_visit_counts = np.concatenate([np.array(c.visit_counts) for c in collectors])
    combined_rewards = np.concatenate([np.array(c.rewards) for c in collectors])

    return ZeroExperienceBuffer(combined_states,
                                combined_visit_counts,
                                combined_rewards)

def LoadExperience(h5file: h5py.File): 
    e = 'experience'
    return ZeroExperienceBuffer(states=np.array(h5file[e]['states']),
                            visit_counts=np.array(h5file[e]['visit_counts']),
                            rewards=np.array(h5file[e]['rewards']))